import pandas as pd
import chardet
import traceback
import os

# 定义目录路径
directory = r"D:\pythonProject\BaiduSyncdisk\LongMoneyTrade\各部分数据\后复权周线"

# 遍历目录下的所有CSV文件
for filename in os.listdir(directory):
    if not filename.endswith('.csv'):
        continue

    file_path = os.path.join(directory, filename)
    print(f"正在处理: {filename}")

    try:
        # 自动检测文件编码
        with open(file_path, 'rb') as f:
            raw_data = f.read()
            result = chardet.detect(raw_data)
            encoding = result['encoding']
            print(f"  检测到文件编码为: {encoding}")

        # 读取 CSV 文件
        df = pd.read_csv(file_path, encoding=encoding, sep=',')

        # 如果只有一列，说明分隔符识别失败，手动处理
        if len(df.columns) == 1:
            header_row = df.columns[0]
            df = pd.read_csv(file_path, encoding=encoding, sep=',', header=None, skiprows=1)
            df.columns = header_row.split(',')

        # 删除已有的'type'和'是否背离'列（如果存在）
        for col in ['type', '是否背离']:
            if col in df.columns:
                df.drop(col, axis=1, inplace=True)
                print(f"  已删除已有列: {col}")

        # 映射实际列名到代码中使用的名称
        column_mapping = {
            '交易结束日期': 'date',
            '周线_后复权最高价': 'high',
            '周线_后复权最低价': 'low',
            'DIF': 'dif'
        }

        # 重命名列
        df = df.rename(columns=column_mapping)

        # 解析日期
        df['date'] = pd.to_datetime(df['date'], errors='coerce')

        # 确保数据按日期排序
        df = df.sort_values('date').dropna(subset=['date'])

        # 初始化两个新列
        df['type'] = ''  # 用于记录高低点类型（peak/trough）
        df['是否背离'] = ''  # 用于记录背离情况

        # 定义窗口大小（前后各10周）
        window_size = 10

        # 使用滚动窗口查找阶段性高点和低点
        peaks_indices = []  # 记录高低点的索引位置

        # 修改：遍历整个数据集
        for i in range(0, len(df)):
            current_dif = df['dif'].iloc[i]

            # 计算前向窗口（前10周）
            start_idx = max(0, i - window_size)
            forward_window = df['dif'].iloc[start_idx:i]

            # 计算后向窗口（后10周，如果不足则取剩余部分）
            end_idx = min(len(df), i + window_size + 1)
            backward_window = df['dif'].iloc[i + 1:end_idx]  # 不包括当前点

            # 判断是否为窗口内的最大值（高点）
            is_peak = False
            if not forward_window.empty and current_dif >= forward_window.max():
                if backward_window.empty or current_dif >= backward_window.max():
                    is_peak = True

            # 判断是否为窗口内的最小值（低点）
            is_trough = False
            if not forward_window.empty and current_dif <= forward_window.min():
                if backward_window.empty or current_dif <= backward_window.min():
                    is_trough = True

            # 标记高低点
            if is_peak:
                df.at[i, 'type'] = 'peak'
                peaks_indices.append(i)
            elif is_trough:
                df.at[i, 'type'] = 'trough'
                peaks_indices.append(i)

        print(f"  找到 {len(peaks_indices)} 个高低点")

        # 如果找到高低点，进行背离检测
        if peaks_indices:
            # 提取所有高低点数据
            peaks_troughs_df = df.iloc[peaks_indices].copy()
            peaks_troughs_df = peaks_troughs_df.sort_values('date')

            # 用于存储背离检测结果
            divergence_results = []

            # 计算每个点的扩展价格（峰点取前后三周最高价，谷点取前后三周最低价）
            extended_prices = []
            for i, row in peaks_troughs_df.iterrows():
                idx = row.name  # 原始索引

                # 计算扩展窗口范围（前后三周）
                start_idx = max(0, idx - 3)
                end_idx = min(len(df) - 1, idx + 3) + 1  # +1 因为切片不包含结束索引

                if row['type'] == 'peak':
                    # 对于峰点，取扩展窗口内的最高价
                    extended_price = df.iloc[start_idx:end_idx]['high'].max()
                else:  # 'trough'
                    # 对于谷点，取扩展窗口内的最低价
                    extended_price = df.iloc[start_idx:end_idx]['low'].min()

                extended_prices.append(extended_price)

            # 添加扩展价格到DataFrame
            peaks_troughs_df['extended_price'] = extended_prices

            # 遍历每个高低点检测背离
            for i in range(len(peaks_troughs_df)):
                current_row = peaks_troughs_df.iloc[i]
                original_index = peaks_troughs_df.index[i]  # 在原始df中的索引位置
                current_date = current_row['date']

                # 顶背离检测 - 只比较前2个高点
                if current_row['type'] == 'peak':
                    # 往前找最多2个peak
                    previous_peaks = []
                    for j in range(i - 1, -1, -1):
                        prev_row = peaks_troughs_df.iloc[j]
                        if prev_row['type'] == 'peak':
                            previous_peaks.append(prev_row)
                            if len(previous_peaks) >= 2:  # 修改：只找2个前高点
                                break

                    # 检查是否存在顶背离
                    for prev_peak in previous_peaks:
                        # 当前DIF值低于前一个峰值，但扩展价格高于前一个峰值的扩展价格
                        if current_row['dif'] < prev_peak['dif'] and current_row['extended_price'] > prev_peak[
                            'extended_price']:
                            divergence_results.append((original_index, '顶背离'))

                            # 输出顶背离比较信息
                            print(f"    顶背离检测: {current_date}")
                            print(
                                f"      当前高点: 扩展最高价={current_row['extended_price']}, DIF={current_row['dif']}")
                            print(f"      前一个高点: 扩展最高价={prev_peak['extended_price']}, DIF={prev_peak['dif']}")
                            print(
                                f"      条件: 价格({current_row['extended_price']} > {prev_peak['extended_price']}), DIF({current_row['dif']} < {prev_peak['dif']})")
                            break

                # 底背离检测 - 只比较前2个低点
                elif current_row['type'] == 'trough':
                    # 往前找最多2个trough
                    previous_troughs = []
                    for j in range(i - 1, -1, -1):
                        prev_row = peaks_troughs_df.iloc[j]
                        if prev_row['type'] == 'trough':
                            previous_troughs.append(prev_row)
                            if len(previous_troughs) >= 2:  # 修改：只找2个前低点
                                break

                    # 检查是否存在底背离
                    for prev_trough in previous_troughs:
                        # 当前DIF值高于前一个谷值，但扩展价格低于前一个谷值的扩展价格
                        if current_row['dif'] > prev_trough['dif'] and current_row['extended_price'] < prev_trough[
                            'extended_price']:
                            divergence_results.append((original_index, '底背离'))

                            # 输出底背离比较信息
                            print(f"    底背离检测: {current_date}")
                            print(
                                f"      当前低点: 扩展最低价={current_row['extended_price']}, DIF={current_row['dif']}")
                            print(
                                f"      前一个低点: 扩展最低价={prev_trough['extended_price']}, DIF={prev_trough['dif']}")
                            print(
                                f"      条件: 价格({current_row['extended_price']} < {prev_trough['extended_price']}), DIF({current_row['dif']} > {prev_trough['dif']})")
                            break

            # 将背离结果写回原始数据框
            for index, divergence_type in divergence_results:
                df.at[index, '是否背离'] = divergence_type
        reverse_mapping = {
            'date': '交易结束日期',
            'high': '周线_后复权最高价',
            'low': '周线_后复权最低价',
            'dif': 'DIF'
        }
        df = df.rename(columns=reverse_mapping)
        # 将结果保存回原始文件（覆盖）
        df.to_csv(file_path, index=False, encoding=encoding)
        print(f"  处理完成: 找到{len(peaks_indices)}个高低点, {len(divergence_results)}个背离信号")
        print("-" * 80)

    except Exception as e:
        print(f"  处理出错: {str(e)}")
        print(traceback.format_exc())
        print("-" * 80)

print("所有文件处理完成！")